import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import tushare as ts
from datetime import datetime, timedelta
from matplotlib import gridspec


class MultiStockTradingStrategy:
    def __init__(self):
        self.pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
        self.stock_data = {}
        self.backtest_results = {}
        self.stock_list = ['002532.SZ', '600519.SH', '000001.SZ']  # 天山铝业、贵州茅台、平安银行

    def fetch_data(self, start_date, end_date, freq='D'):
        """获取多只股票数据"""
        print(f"正在获取{freq}线数据...")

        for stock in self.stock_list:
            if freq == 'D':  # 日线
                df = self.pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
            elif freq == 'W':  # 周线
                df = self.pro.weekly(ts_code=stock, start_date=start_date, end_date=end_date)
            elif freq == '60min':  # 60分钟线
                df = self.pro.pro_bar(ts_code=stock, freq='60min', start_date=start_date, end_date=end_date)
                df.rename(columns={'trade_time': 'trade_date'}, inplace=True)

            if len(df) == 0:
                continue

            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df = df.sort_values('trade_date')

            # 计算技术指标
            df['ma5'] = df['close'].rolling(5).mean()
            df['ma20'] = df['close'].rolling(20).mean()
            df['rsi'] = self._calculate_rsi(df)
            df['atr'] = self._calculate_atr(df)

            self.stock_data[stock] = df
            print(f"{stock} {freq}线数据获取完成，共{len(df)}条记录")

    def _calculate_rsi(self, df, window=14):
        """计算RSI指标"""
        delta = df['close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)

        avg_gain = gain.rolling(window).mean()
        avg_loss = loss.rolling(window).mean()

        rs = avg_gain / avg_loss
        return 100 - (100 / (1 + rs))

    def _calculate_atr(self, df, window=14):
        """计算ATR(平均真实波幅)"""
        high_low = df['high'] - df['low']
        high_close = (df['high'] - df['close'].shift()).abs()
        low_close = (df['low'] - df['close'].shift()).abs()

        tr = pd.concat([high_low, high_close, low_close], axis=1).max(axis=1)
        return tr.rolling(window).mean()

    def generate_signals(self, stock):
        """生成交易信号"""
        df = self.stock_data[stock].copy()

        # 双均线策略 + RSI + 突破策略
        df['signal'] = 0

        # 买入信号条件
        buy_condition = (
                (df['ma5'] > df['ma20']) &
                (df['rsi'] > 30) & (df['rsi'] < 70) &
                (df['close'] > df['high'].rolling(5).max().shift(1))
        )

        # 卖出信号条件
        sell_condition = (
                (df['ma5'] < df['ma20']) |
                (df['rsi'] > 70) |
                (df['close'] < df['low'].rolling(5).min().shift(1))
        )

        df.loc[buy_condition, 'signal'] = 1
        df.loc[sell_condition, 'signal'] = -1

        self.stock_data[stock] = df
        return df

    def backtest(self, stock, initial_capital=1000000):
        """回测单只股票策略"""
        df = self.stock_data[stock].copy()
        capital = initial_capital
        positions = 0
        entry_price = 0
        trade_log = []
        portfolio_values = []

        for i, row in df.iterrows():
            date = row['trade_date']
            price = row['close']
            atr = row['atr']
            signal = row['signal']

            # 卖出逻辑
            if signal == -1 and positions > 0:
                profit = (price - entry_price) * positions
                capital += price * positions
                trade_log.append({
                    'date': date, 'action': 'sell', 'price': price,
                    'shares': positions, 'profit': profit,
                    'capital': capital, 'portfolio': capital
                })
                positions = 0

            # 买入逻辑
            elif signal == 1 and positions == 0 and not pd.isna(atr):
                risk_per_trade = 0.02  # 2%风险
                shares = int((capital * risk_per_trade) / (1.5 * atr))
                if shares > 0 and capital >= price * shares:
                    capital -= price * shares
                    positions = shares
                    entry_price = price
                    trade_log.append({
                        'date': date, 'action': 'buy', 'price': price,
                        'shares': shares, 'capital': capital,
                        'portfolio': capital + price * shares,
                        'stop_loss': price - 1.5 * atr
                    })

            # 每日记录组合价值
            portfolio_value = capital + (positions * price if positions > 0 else 0)
            portfolio_values.append({
                'date': date,
                'value': portfolio_value,
                'price': price
            })

        # 计算收益率
        initial_price = df.iloc[0]['close']
        final_price = df.iloc[-1]['close']
        buy_hold_return = (final_price / initial_price - 1) * 100

        strategy_return = (portfolio_values[-1]['value'] / initial_capital - 1) * 100

        result = {
            'trade_log': pd.DataFrame(trade_log),
            'portfolio': pd.DataFrame(portfolio_values),
            'buy_hold_return': buy_hold_return,
            'strategy_return': strategy_return
        }

        self.backtest_results[stock] = result
        return result

    def run_all_backtests(self, start_date, end_date, freq='D'):
        """运行所有股票的回测"""
        self.fetch_data(start_date, end_date, freq)

        for stock in self.stock_list:
            if stock in self.stock_data:
                print(f"\n正在回测 {stock}...")
                self.generate_signals(stock)
                self.backtest(stock)

                # 打印简要结果
                result = self.backtest_results[stock]
                print(f"策略收益率: {result['strategy_return']:.2f}%")
                print(f"买入持有收益率: {result['buy_hold_return']:.2f}%")
                print(f"交易次数: {len(result['trade_log']) // 2}")

    def visualize_results(self):
        """可视化回测结果"""
        if not self.backtest_results:
            print("请先运行回测")
            return

        plt.figure(figsize=(16, 12))
        gs = gridspec.GridSpec(3, 2, height_ratios=[2, 1, 1])

        # 收益率曲线比较
        ax1 = plt.subplot(gs[0, :])
        for stock, result in self.backtest_results.items():
            portfolio = result['portfolio']
            ax1.plot(portfolio['date'], portfolio['value'] / 1000000,
                     label=f"{stock} 策略净值", linewidth=2)
            ax1.plot(portfolio['date'], portfolio['price'] / portfolio['price'].iloc[0],
                     '--', label=f"{stock} 股价走势", alpha=0.7)
        ax1.set_title('多股票策略净值比较 (标准化)', fontsize=15)
        ax1.set_ylabel('净值倍数', fontsize=12)
        ax1.legend(fontsize=10)
        ax1.grid(True, alpha=0.3)

        # 各股票K线图+信号
        for i, stock in enumerate(self.stock_list[:2]):
            if stock not in self.backtest_results:
                continue

            ax = plt.subplot(gs[1 + i, :])
            df = self.stock_data[stock]
            portfolio = self.backtest_results[stock]['portfolio']
            trade_log = self.backtest_results[stock]['trade_log']

            # 绘制K线
            ax.plot(df['trade_date'], df['close'], label='收盘价', color='blue', alpha=0.7)

            # 标记买卖点
            buys = trade_log[trade_log['action'] == 'buy']
            sells = trade_log[trade_log['action'] == 'sell']
            ax.scatter(buys['date'], buys['price'], color='green',
                       marker='^', s=100, label='买入')
            ax.scatter(sells['date'], sells['price'], color='red',
                       marker='v', s=100, label='卖出')

            ax.set_title(f'{stock} K线及交易信号', fontsize=12)
            ax.legend()
            ax.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()

        # 单独显示每只股票的交易信号细节
        for stock in self.stock_list:
            if stock not in self.backtest_results:
                continue

            self._plot_stock_details(stock)

    def _plot_stock_details(self, stock):
        """绘制单只股票的详细交易信号"""
        df = self.stock_data[stock]
        result = self.backtest_results[stock]

        plt.figure(figsize=(16, 10))
        gs = gridspec.GridSpec(3, 1, height_ratios=[2, 1, 1])

        # K线和技术指标
        ax1 = plt.subplot(gs[0])
        ax1.plot(df['trade_date'], df['close'], label='收盘价', color='blue')
        ax1.plot(df['trade_date'], df['ma5'], label='5日均线', color='orange', alpha=0.7)
        ax1.plot(df['trade_date'], df['ma20'], label='20日均线', color='green', alpha=0.7)

        # 标记买卖点
        trades = result['trade_log']
        buys = trades[trades['action'] == 'buy']
        sells = trades[trades['action'] == 'sell']
        ax1.scatter(buys['date'], buys['price'], color='green',
                    marker='^', s=100, label='买入')
        ax1.scatter(sells['date'], sells['price'], color='red',
                    marker='v', s=100, label='卖出')

        ax1.set_title(f'{stock} - 交易信号详情', fontsize=15)
        ax1.legend()
        ax1.grid(True, alpha=0.3)

        # RSI指标
        ax2 = plt.subplot(gs[1])
        ax2.plot(df['trade_date'], df['rsi'], label='RSI', color='purple')
        ax2.axhline(30, linestyle='--', color='green', alpha=0.5)
        ax2.axhline(70, linestyle='--', color='red', alpha=0.5)
        ax2.set_ylabel('RSI')
        ax2.grid(True, alpha=0.3)

        # 成交量
        ax3 = plt.subplot(gs[2])
        ax3.bar(df['trade_date'], df['vol'], color='gray', alpha=0.7)
        ax3.set_ylabel('成交量')
        ax3.grid(True, alpha=0.3)

        plt.tight_layout()
        plt.show()


# 使用示例
if __name__ == "__main__":
    strategy = MultiStockTradingStrategy()

    # 设置回测日期范围
    end_date = datetime.now().strftime('%Y%m%d')
    start_date = (datetime.now() - timedelta(days=365)).strftime('%Y%m%d')  # 1年数据

    # 运行回测 (可更改freq为'D'(日线), 'W'(周线)或'60min'(60分钟线))
    strategy.run_all_backtests(start_date, end_date, freq='D')

    # 可视化结果
    strategy.visualize_results()